﻿#include <iostream>
#include <thread>
#include <mutex>
#include <arpa/inet.h>
#include <unistd.h>
#include<string.h>
#include<sys/epoll.h>
#include<fcntl.h>

using namespace std;

#define PORT 9999
#define MAX_CLIENTS 10
#define BUFFER_SIZE 1024
#define MAX_EVENTS 1024
#define MAX_FILENAME_LEN 256
#define HEADER_SIZE 5 //消息头部大小：4字节长度+1字节类型
//定义消息类型
#define MSG_TEXT 'T'
#define FILE_UPLOAD 'U'
#define FILE_DOWNLOAD 'D'

typedef struct {
    int fd;
    FILE* fp;
}Client;
  
Client clients[MAX_EVENTS]; //客户端数组
int clnt_count = 0;

//处理客户端事件
void handle_client_event(Client* client, epoll_event* event); 
//广播文件名给所有客户端
void broadcast_file_name(const char* file_name, int sender_fd);
//解析消息头部
void parse_header(const char* buffer, size_t* length, char* type);
//检查文件名是否合法
int is_valid_name(const char* file_name);
void send_file(const char *file_name, int client_fd);
//void set_non_blocking(int fd);



int main() {
    //服务器端建立监听
    int serv_sock, clnt_sock;
    struct sockaddr_in serv_addr, clnt_addr;
    socklen_t clnt_addr_len;
    // 创建套接字
    serv_sock = socket(AF_INET, SOCK_STREAM, 0);
    if (serv_sock == 0) {
        perror("socket error");
        return -1;
    }
    //set_non_blocking(serv_sock);
    //设置端口复用：允许套接字在关闭后可以立马重用已绑定的地址和端口
    int opt = 1;
    setsockopt(serv_sock, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    
    //绑定端口
    memset(&serv_addr, 0, sizeof(serv_addr));
    serv_addr.sin_family = AF_INET;
    serv_addr.sin_addr.s_addr = htonl(INADDR_ANY);
    serv_addr.sin_port = htons(PORT);
    if (bind(serv_sock, (struct sockaddr*)&serv_addr, sizeof(serv_addr)) < 0) {
        perror("bind error");
        close(serv_sock);
        return -1;
    }
    //监听
    if (listen(serv_sock, MAX_CLIENTS) < 0) {
        perror("listen error");
        close(serv_sock);
        return -1;
    }
    cout << "Server listening on port: " << PORT << endl;

    int epfd, event_cnt;
    epoll_event event;
    epoll_event* all_events = new epoll_event[MAX_EVENTS]; //已发生事件数组
    //创建一个epoll实例
    epfd = epoll_create(1);
    if (epfd == -1) {
        perror("epoll_create error");
        close(serv_sock);
        return -1;
    }
    //向epoll实例注册一个serv_sock描述符，并将监听的事件类型设置为可读状态
    //当客户端尝试连接服务器，serv_sock会变为可读状态，从而被epoll_wait检测到，并保存在事件结构体中返回
    event.events = EPOLLIN; //监控的事件类型为可读
    event.data.fd = serv_sock;
    epoll_ctl(epfd, EPOLL_CTL_ADD, serv_sock, &event);
    while (1) {
        event_cnt = epoll_wait(epfd, all_events, MAX_EVENTS, 1000);
        if (event_cnt == -1) {
            perror("epoll_wait error");
            break;
        }
        if (event_cnt == 0) continue;
        //处理事件
        for (int i = 0; i < event_cnt; ++i) {
            if (all_events[i].data.fd == serv_sock) {  //处理新连接
                clnt_sock = accept(serv_sock, (sockaddr*)&clnt_addr, &clnt_addr_len);
                //set_non_blocking(clnt_sock);
                printf("new connection: fd=%d,ip=%s,port=%d\n", clnt_sock, inet_ntoa(clnt_addr.sin_addr), ntohs(clnt_addr.sin_port));
                clnt_count++;
                //在epoll实例中注册客户端套接字描述符
                event.events = EPOLLIN ; //监控的事件类型为读事件，模式为水平触发
                event.data.fd = clnt_sock;
                epoll_ctl(epfd, EPOLL_CTL_ADD, clnt_sock, &event);
                //将新连接的客户端添加到客户端列表中
                for (int j = 0; j < MAX_EVENTS; ++j) {
                    if (clients[j].fd == 0) {
                        clients[j].fd = clnt_sock;
                        clients[j].fp = NULL;
                        break;
                    }
                }
            }
            else { //处理客户端事件
                for (int j = 0; j < MAX_EVENTS; ++j) {
                    if (clients[j].fd == all_events[i].data.fd) {
                        handle_client_event(&clients[j], &all_events[i]);
                        break;
                    }
                }
            }
        }
    }
    close(serv_sock);
    close(epfd);
    return 0;
    
}

void handle_client_event(Client* client, epoll_event* event){
    if (event->events & EPOLLIN) {
        char buffer[BUFFER_SIZE] = "";
        ssize_t len_read;
        //先读一个buffer的数据
        len_read = read(client->fd, buffer, sizeof(buffer));
        if (len_read == 0) { //客户端断开
            printf("client %d disconnected\n", client->fd);
            close(client->fd);
            memset(client, 0, sizeof(Client));
            cout << "client counts:" << --clnt_count << endl;
        }
        //处理buffer中的数据
        char packet_type;
        size_t length;
        char file_name[MAX_FILENAME_LEN]; //存放文件名的数组
        memset(file_name, 0, sizeof(file_name));
        //解析消息头部
        if (len_read < HEADER_SIZE) { //数据不足以解析头部
            fprintf(stderr, "Invalid header from client %d\n", client->fd);
            return;
        }
        parse_header(buffer, &length, &packet_type);
        //处理消息正文
        char* content = buffer + HEADER_SIZE; //消息正文的起始位置
        if (packet_type == MSG_TEXT) { //文本消息
            //检查是否收到完整的文本,数据不足时返回
            if (len_read < HEADER_SIZE + length) {
                printf("text is incomplete!\n");
                return;
            }
            for (int i = 0; i < MAX_EVENTS; ++i) { //广播给其他客户端
                if (clients[i].fd != 0 && clients[i].fd != client->fd)
                    send(clients[i].fd, buffer, HEADER_SIZE + length, 0);
            }
            printf("BroadCast message from client: %d\n", client->fd);
        }
        else if (packet_type == FILE_DOWNLOAD) { //文件下载请求
            //提取文件名
            strncpy(file_name, content, length);
            if (is_valid_name(file_name))
                send_file(file_name, client->fd);
            else
                fprintf(stderr, "Invalid file name request from client %d\n", client->fd);
        }
        else if (packet_type == FILE_UPLOAD) { //客户端上传文件
            printf("------FILE_UPLOAD------\n");
            //获取文件名
            const char* FileNameStart = content;
            const char* FileNameEnd = strchr(FileNameStart, '\0');
            size_t FileNameLen = FileNameEnd - FileNameStart;
            strncpy(file_name, FileNameStart, FileNameLen);
            file_name[FileNameLen] = 0;
            printf("------File name: %s------\n", file_name);
            if (is_valid_name(file_name)) { //文件名正确
                //创建一个空文件
                client->fp = fopen(file_name, "wb");
                if (!client->fp) {
                    fprintf(stderr, "Error opening file for: %s\n", file_name);
                }
                //循环接收文件
                int file_recved = 0; //目前接收
                int file_remaining = length; //剩余需要接收的
                printf("------Start File recving.------\n");
                while (1) {
                    memset(buffer, 0, sizeof(buffer));
                    len_read = read(client->fd, buffer, sizeof(buffer));
                    printf("------len_read:%d,file_recved:%d,file_remaining:%d------\n", len_read, file_recved, file_remaining);
                    fwrite(buffer, 1, len_read, client->fp);  //将接收到的文件内容写入本地文件
                    if (len_read == file_remaining) { //读完文件
                        file_recved += len_read;
                        file_remaining -= len_read;
                        printf("------len_read:%d,file_recved:%d,file_remaining:%d------\n", len_read, file_recved, file_remaining);
                        break;
                    }
                    else if (len_read > 0) { //还有剩余的文件数据没有传过来
                        file_recved += len_read;
                        file_remaining -= len_read;
                        continue;
                    }
                    else if (len_read <= 0) {
                        continue;
                    }
                }
                if (file_recved >= length) { //文件传输完成
                    //关闭文件
                    fclose(client->fp);
                    client->fp = NULL;
                    printf("------File: %s upload complete!------\n", file_name);
                    //广播文件名给其他客户端
                    broadcast_file_name(file_name, client->fd);
                }
            }
            else {
                fprintf(stderr, "Invalid file name from client %d\n", client->fd);
            }
        }
    }
}


void broadcast_file_name(const char* file_name, int sender_fd) {
    char message[MAX_FILENAME_LEN + HEADER_SIZE];
    size_t message_len = strlen(file_name) + 1;
    *(uint32_t*)message = htonl(message_len); //消息长度
    message[4] = FILE_UPLOAD; //消息类型
    strcpy(message + HEADER_SIZE, file_name); //消息内容
    //广播消息
    for (int i = 0; i < MAX_EVENTS; ++i) {
        if (clients[i].fd != 0 )
            send(clients[i].fd, message, HEADER_SIZE + message_len, 0);
    }
    printf("------Broadcast file name: %s------\n", file_name);
}

void parse_header(const char* buffer, size_t* length, char* type){
    *length = ntohl(*(uint32_t*)buffer); //解析长度字段:先转换为unsigned int型指针，然后解引用得到里面的值
    *type = buffer[4]; //解析消息类型
    return;
}

int is_valid_name(const char* file_name){
    for (int i = 0; file_name[i] != 0; ++i) {
        if (file_name[i] == '/' || file_name[i] == '\\') return 0; //文件名包含非法字符
    }
    return 1;
}

void send_file(const char* file_name, int client_fd){
    FILE* fp = fopen(file_name, "rb");
    if (!fp) {
        fprintf(stderr, "Error opening file: %s\n", file_name);
        return;
    }
    //获取文件大小
    fseek(fp, 0,SEEK_END);
    size_t file_size = ftell(fp);
    fseek(fp, 0, SEEK_SET);
    //发送文件头部
    char header[HEADER_SIZE];
    *(uint32_t*)header = htonl(file_size); //文件大小
    header[4] = FILE_DOWNLOAD; //消息类型
    send(client_fd, header, HEADER_SIZE, 0);
    //分块发送文件数据
    char send_buffer[BUFFER_SIZE];
    memset(send_buffer, 0, BUFFER_SIZE);
    size_t len_read;
    ssize_t total_byteSent = 0;
    while ((len_read = fread(send_buffer,1,BUFFER_SIZE,fp)) > 0){
        ssize_t byteSent = send(client_fd, send_buffer, len_read, 0);
        if (byteSent == -1) {
            perror("Error sending");
            break;
        }
        total_byteSent += byteSent;
        printf("Sent: %d bytes, total sent: %d bytes, file size: %d\n", byteSent, total_byteSent, file_size);
    }
    fclose(fp);
    printf("File %s sent to client %d\n", file_name, client_fd);
}
